"""
交错网格模块：用于两相流计算的交错网格系统

该模块实现了用于两相流计算的交错网格系统，包括：
1. 网格数据的存储和管理
2. 物理量的更新和插值
3. 源项的计算和管理
4. 边界条件的处理
"""

import numpy as np
from typing import Dict, Optional, Tuple, List
from dataclasses import dataclass
from enum import Enum, auto
import CoolProp.CoolProp as CP

class BoundaryType(Enum):
    """边界条件类型"""
    INLET = auto()          # 入口边界
    OUTLET = auto()         # 出口边界
    ROD_WALL = auto()       # 燃料棒壁面
    SHROUD_WALL = auto()    # 包壳壁面
    CHANNEL_CONNECTION = auto()  # 子通道连接面

@dataclass
class BoundaryCondition:
    """边界条件数据结构"""
    type: BoundaryType
    location: Dict[str, int]  # 边界位置信息（如rod_id, channel_id等）
    values: Dict[str, float]  # 边界值字典

@dataclass
class GridData:
    """网格数据结构"""
    # 标量场（定义在网格中心）
    scalar_fields: Dict[str, np.ndarray] = None
    # 轴向速度场（定义在x方向网格面上）
    axial_velocities: Dict[str, np.ndarray] = None
    # 横向速度场（定义在y方向网格面上）
    lateral_velocities: Dict[str, np.ndarray] = None
    # 源项（定义在网格中心）
    source_terms: Dict[str, np.ndarray] = None
    # 上一时间步的数据
    old_data: Dict[str, np.ndarray] = None

class StaggeredGrid:
    """交错网格类,用于管理网格数据、物理参数和边界条件"""
    
    def __init__(self, config: dict):
        """
        初始化交错网格
        
        参数:
            config (dict): 配置参数字典
        """
        # 网格参数
        self.nx = config['geometry']['n_axial']
        self.ny = config['geometry']['n_channels']
        self.dx = config['geometry']['channel_length'] / self.nx
        self.dy = config['geometry']['rod_pitch']
        self.dt = config['numerical']['time_step']
        
        # 初始化标量场(定义在网格中心)
        self.scalar_fields = {
            'pressure': np.zeros((self.nx, self.ny)),
            'temperature': np.zeros((self.nx, self.ny)),
            'density': np.zeros((self.nx, self.ny)),
            'enthalpy': np.zeros((self.nx, self.ny)),
            'void_fraction': np.zeros((self.nx, self.ny)),
            'density_liquid': np.zeros((self.nx, self.ny)),
            'density_vapor': np.zeros((self.nx, self.ny)),
            'viscosity_liquid': np.zeros((self.nx, self.ny)),
            'viscosity_vapor': np.zeros((self.nx, self.ny)),
            'conductivity_liquid': np.zeros((self.nx, self.ny)),
            'conductivity_vapor': np.zeros((self.nx, self.ny))
        }
        
        # 初始化速度场(定义在网格界面)
        self.velocity_fields = {
            'axial': np.zeros((self.nx+1, self.ny)),    # u速度,定义在x方向网格界面
            'lateral': np.zeros((self.nx, self.ny+1))    # v速度,定义在y方向网格界面
        }
        
        # 初始化源项
        self.source_terms = {}
        
        # 存储上一时间步的数据
        self.old_scalar_fields = {}
        self.old_velocity_fields = {}
        
        # 初始化流体物性参数
        self.fluid_properties = {
            'reference_pressure': config['physics']['reference_pressure'],
            'reference_temperature': config['physics']['reference_temperature'],
            'critical_pressure': config['physics']['critical_pressure'],
            'critical_temperature': config['physics']['critical_temperature'],
            'liquid_density_ref': config['physics']['liquid_density_ref'],
            'liquid_viscosity_ref': config['physics']['liquid_viscosity_ref'],
            'liquid_conductivity_ref': config['physics']['liquid_conductivity_ref'],
            'liquid_cp_ref': config['physics']['liquid_cp_ref'],
            'vapor_density_ref': config['physics']['vapor_density_ref'],
            'vapor_viscosity_ref': config['physics']['vapor_viscosity_ref'],
            'vapor_conductivity_ref': config['physics']['vapor_conductivity_ref'],
            'vapor_cp_ref': config['physics']['vapor_cp_ref'],
            'latent_heat': config['physics']['latent_heat'],
            'surface_tension_ref': config['physics']['surface_tension_ref']
        }
        
        # 初始化边界条件
        self.boundary_conditions = {
            'inlet_pressure': config['boundary']['inlet_pressure'],
            'inlet_temperature': config['boundary']['inlet_temperature'],
            'inlet_mass_flux': config['boundary']['inlet_mass_flux'],
            'inlet_quality': config['boundary']['inlet_quality'],
            'outlet_pressure': config['boundary']['outlet_pressure'],
            'wall_heat_flux': np.zeros((self.nx, self.ny)),
            'wall_temperature': np.zeros((self.nx, self.ny))
        }
        
        # 初始化物理模型参数
        self.model_parameters = {
            'drift_flux_c0': config['models']['drift_flux_c0'],
            'drift_flux_vgj': config['models']['drift_flux_vgj'],
            'single_phase_htc_coef': config['models']['single_phase_htc_coef'],
            'nucleate_boiling_coef': config['models']['nucleate_boiling_coef'],
            'friction_coef_laminar': config['models']['friction_coef_laminar'],
            'friction_coef_turbulent': config['models']['friction_coef_turbulent'],
            'turbulent_mixing_coef': config['models']['turbulent_mixing_coef']
        }
        
    def update_fluid_properties(self):
        """更新流体物性参数"""
        # 获取当前温度和压力场
        pressure = self.get_scalar_field('pressure')
        temperature = self.get_scalar_field('temperature')
        
        # 使用CoolProp计算物性
        for i in range(self.nx):
            for j in range(self.ny):
                P = pressure[i,j]
                T = temperature[i,j]
                
                # 液相物性
                self.scalar_fields['density_liquid'][i,j] = CP.PropsSI('D', 'P', P, 'T', T, 'water')
                self.scalar_fields['viscosity_liquid'][i,j] = CP.PropsSI('V', 'P', P, 'T', T, 'water')
                self.scalar_fields['conductivity_liquid'][i,j] = CP.PropsSI('L', 'P', P, 'T', T, 'water')
                
                # 气相物性
                self.scalar_fields['density_vapor'][i,j] = CP.PropsSI('D', 'P', P, 'Q', 1, 'water')
                self.scalar_fields['viscosity_vapor'][i,j] = CP.PropsSI('V', 'P', P, 'Q', 1, 'water')
                self.scalar_fields['conductivity_vapor'][i,j] = CP.PropsSI('L', 'P', P, 'Q', 1, 'water')
                
    def update_wall_conditions(self, heat_flux: np.ndarray):
        """
        更新壁面边界条件
        
        参数:
            heat_flux (np.ndarray): 壁面热流密度分布
        """
        if heat_flux.shape == self.boundary_conditions['wall_heat_flux'].shape:
            self.boundary_conditions['wall_heat_flux'] = heat_flux
        else:
            raise ValueError("热流密度数组形状不匹配")
            
    def get_fluid_property(self, property_name: str) -> float:
        """获取流体物性参数"""
        if property_name in self.fluid_properties:
            return self.fluid_properties[property_name]
        else:
            raise KeyError(f"物性参数 {property_name} 不存在")
            
    def get_model_parameter(self, parameter_name: str) -> float:
        """获取物理模型参数"""
        if parameter_name in self.model_parameters:
            return self.model_parameters[parameter_name]
        else:
            raise KeyError(f"模型参数 {parameter_name} 不存在")
            
    def get_boundary_condition(self, condition_name: str) -> float:
        """获取边界条件参数"""
        if condition_name in self.boundary_conditions:
            return self.boundary_conditions[condition_name]
        else:
            raise KeyError(f"边界条件 {condition_name} 不存在")
            
    def store_current_step(self):
        """存储当前时间步的数据"""
        self.old_scalar_fields = {
            key: value.copy() for key, value in self.scalar_fields.items()
        }
        self.old_velocity_fields = {
            key: value.copy() for key, value in self.velocity_fields.items()
        }
        
    def get_scalar_field(self, name: str) -> np.ndarray:
        """获取标量场"""
        return self.scalar_fields[name]
        
    def get_old_scalar_field(self, name: str) -> np.ndarray:
        """获取上一时间步的标量场"""
        return self.old_scalar_fields[name]
        
    def get_axial_velocity_field(self) -> np.ndarray:
        """获取轴向速度场"""
        return self.velocity_fields['axial']
        
    def get_lateral_velocity_field(self) -> np.ndarray:
        """获取横向速度场"""
        return self.velocity_fields['lateral']
        
    def get_old_axial_velocity_field(self) -> np.ndarray:
        """获取上一时间步的轴向速度场"""
        return self.old_velocity_fields['axial']
        
    def get_old_lateral_velocity_field(self) -> np.ndarray:
        """获取上一时间步的横向速度场"""
        return self.old_velocity_fields['lateral']
        
    def get_source_term(self, name: str) -> np.ndarray:
        """获取源项"""
        return self.source_terms.get(name, np.zeros((self.nx, self.ny)))
        
    def set_source_term(self, name: str, value: np.ndarray):
        """设置源项"""
        self.source_terms[name] = value
        
    def update_variables(self, residuals: Dict[str, Dict[str, np.ndarray]], dt: float):
        """
        基于残差更新物理量
        
        参数:
            residuals (Dict[str, Dict[str, np.ndarray]]): 各守恒方程的残差
            dt (float): 时间步长
        """
        # 更新密度
        rho = self.get_scalar_field('density')
        rho_residual = residuals['mass']
        rho -= dt * rho_residual
        
        # 更新速度
        u = self.get_axial_velocity_field()
        v = self.get_lateral_velocity_field()
        u_residual = residuals['momentum_x']
        v_residual = residuals['momentum_y']
        u -= dt * u_residual
        v -= dt * v_residual
        
        # 更新焓
        h = self.get_scalar_field('enthalpy')
        h_residual = residuals['energy']
        h -= dt * h_residual
        
        # 更新压力
        # 这里使用简化的SIMPLE算法更新压力
        p = self.get_scalar_field('pressure')
        p_correction = self._solve_pressure_correction()
        p += p_correction
        
        # 更新温度和空泡率(通过状态方程)
        self._update_state_variables()
        
        # 应用边界条件
        self.apply_boundary_conditions()
        
    def apply_boundary_conditions(self):
        """应用边界条件"""
        self._apply_inlet_boundary()
        self._apply_outlet_boundary()
        self._apply_wall_boundary()
        self._apply_symmetry_boundary()
        
    def _apply_inlet_boundary(self):
        """应用入口边界条件"""
        # 设置入口速度
        self.velocity_fields['axial'][0,:] = self.inlet_velocity
        
        # 设置入口温度
        self.scalar_fields['temperature'][0,:] = self.inlet_temperature
        
        # 设置入口空泡率
        self.scalar_fields['void_fraction'][0,:] = self.inlet_void_fraction
        
    def _apply_outlet_boundary(self):
        """应用出口边界条件"""
        # 设置出口压力
        self.scalar_fields['pressure'][-1,:] = self.outlet_pressure
        
        # 零梯度条件
        self.scalar_fields['temperature'][-1,:] = self.scalar_fields['temperature'][-2,:]
        self.scalar_fields['void_fraction'][-1,:] = self.scalar_fields['void_fraction'][-2,:]
        self.velocity_fields['axial'][-1,:] = self.velocity_fields['axial'][-2,:]
        
    def _apply_wall_boundary(self):
        """应用壁面边界条件"""
        # 燃料棒壁面无滑移条件
        self.velocity_fields['axial'][:,0] = 0.0
        self.velocity_fields['axial'][:,-1] = 0.0
        self.velocity_fields['lateral'][:,0] = 0.0
        self.velocity_fields['lateral'][:,-1] = 0.0
        
        # 壁面热流密度条件(如果指定)
        if hasattr(self, 'wall_heat_flux'):
            self._apply_wall_heat_flux()
            
    def _apply_symmetry_boundary(self):
        """应用对称边界条件"""
        # 子通道间界面的对称条件
        # 标量量对称
        for field in self.scalar_fields.values():
            field[:,0] = field[:,1]
            field[:,-1] = field[:,-2]
            
        # 速度场反对称
        self.velocity_fields['lateral'][:,0] = -self.velocity_fields['lateral'][:,1]
        self.velocity_fields['lateral'][:,-1] = -self.velocity_fields['lateral'][:,-2]
        
    def _solve_pressure_correction(self) -> np.ndarray:
        """
        求解压力校正方程
        
        返回:
            np.ndarray: 压力校正量
        """
        # 这里使用简化的SIMPLE算法
        # 实际应用中需要求解压力泊松方程
        return np.zeros((self.nx, self.ny))
        
    def _update_state_variables(self):
        """更新状态变量(温度和空泡率)"""
        # 通过状态方程更新温度
        p = self.scalar_fields['pressure']
        h = self.scalar_fields['enthalpy']
        self.scalar_fields['temperature'] = self._calculate_temperature(p, h)
        
        # 更新空泡率
        self.scalar_fields['void_fraction'] = self._calculate_void_fraction()
        
    def _calculate_temperature(self, p: np.ndarray, h: np.ndarray) -> np.ndarray:
        """
        计算温度(通过状态方程)
        
        参数:
            p (np.ndarray): 压力场
            h (np.ndarray): 焓值场
            
        返回:
            np.ndarray: 温度场
        """
        return self.heat_transfer.calculate_temperature(p, h)
        
    def _calculate_void_fraction(self) -> np.ndarray:
        """
        计算空泡率(通过相变模型)
        
        返回:
            np.ndarray: 空泡率场
        """
        p = self.scalar_fields['pressure']
        h = self.scalar_fields['enthalpy']
        return self.heat_transfer.calculate_void_fraction(p, h)
        
    def _apply_wall_heat_flux(self):
        """应用壁面热流密度边界条件"""
        # 在壁面处施加热流密度
        pass
        
    def get_dimensions(self) -> Tuple[int, int]:
        """获取网格的维度"""
        return self.nx, self.ny
        